﻿#'
#' @TODO  组间HALLMARK富集分数差异
#' @title ## 组间HALLMARK富集分数差异
#' @description 分组间ssgsea 通路富集得分比较，秩和检验判断组间得分显著性。
#' @param exp *data.frame* 标准化后的表达谱 基因在行 样本在列的
#' @param group_infor 样本分组信息，仅用于划分exp中的表达谱样本是有风险评分的.包含`sample`与`group`两列，且表头一致
#' @param score_res 得分计算结果，如果为NULL则程序自动根据`path_way`计算，如果不是NULL，程序跳过计算步骤，自动利用得分结果进行绘图，在此项存在的情况下，`path_way`可以是NULL
#' @param path_way 通路名称，单一字符串，模糊匹配。或者是很多具体的、唯一匹配的通路名称，字符串向量
#' @param col_used 自定义聚类注释颜色，否则使用set1 配色。注：配色顺序为`group_infor$group`的出场顺序。(字母表顺序)
#' @param saveplot 是否保存图片
#' @param output_dir 结果输出路径 ，需要以`/`结尾
#' @param display_all 是否展示全部通路结果，而不是只展示显著的,默认是F，只展示显著的通路结果
#' @param cluster_name 分组信息表头，如果为NULL则使用group
#' @param h 图高
#' @param w 图宽
#' @examples
#' source("/home/innertech/RCodes/Project_wangyk/Codelib_YK/before_server/pathway_ssgsea_score.R")
#'
#' # 高低风险组间
#' res1 <- score_compute_heatmap(
#'     exp = exp_normalized, group_infor = training[[3]] %>%
#'         select(sample, riskgroup) %>% rename(group = riskgroup),
#'     score_res = NULL,
#'     path_way = "hallmark", output_dir = output_dir,
#'     saveplot = T, col_used = NULL, var_name = "coadtest", display_all = T
#' )
#'
#' # 聚类三分组间
#' res2 <- score_compute_heatmap(
#'     exp = TCGA_TPM[, a], group_infor = group_infor,
#'     score_res = as.data.frame(rs_score), path_way = reactome, output_dir = paste0(getwd(), "/"),
#'     saveplot = T, col_used = c("#FF9288", "#00D55C", "#82B7FF"), var_name = "RS_used", display_all = T
#' )
#'
#' @return *list*
#' @examples
#' @author *WYK*
#'
score_compute_heatmap <- function(exp = NULL, group_infor = NULL, score_res = NULL,
                                  path_way = "hallmark", output_dir = "./",
                                  saveplot = T, col_used = NULL, var_name = NULL, cluster_name = NULL, display_all = F,
                                  h = 5, w = 5) {
    set.seed(1110)
    if (is.null(var_name)) {
        var_name <- paste0(sample(letters, 4), collapse = "")
    }

    if (is.null(score_res)) {
        group_infor <- group_infor[{
            match(colnames(exp), group_infor$sample) %>%
                na.omit() %>%
                as.numeric()
        }, ]
        exp <- exp[, {
            match(group_infor$sample, colnames(exp)) %>%
                na.omit() %>%
                as.numeric()
        }]
        exp_normalized <- exp
    }

    library(cowplot)
    library(tidyverse)
    library(RColorBrewer)

    if (length(score_res) == 0) {
        pathway <- getMsigDBGeneLst(path_way)
        names(pathway) <- str_remove_all(names(pathway), "HALLMARK_")
        names(pathway) <- str_remove_all(names(pathway), "KEGG_") %>% str_replace_all(., "_", " ")

        ssgsea_score <- GSVA::gsva(
            expr = exp_normalized[, group_infor$sample] %>% as.matrix(),
            gset.idx.list = pathway,
            method = "ssgsea",
            ssgsea.norm = TRUE,
            verbose = T,
            parallel.sz = 50
        ) %>% as.data.frame() # signature 'matrix,list'

        colnames(ssgsea_score) <- str_replace_all(colnames(ssgsea_score), "_", " ")
    } else {
        ssgsea_score <- score_res
        colnames(ssgsea_score) <- str_replace_all(colnames(ssgsea_score), "_", " ")
    }

    data_forboxplot <- ssgsea_score %>% as.data.frame() %>% 
        tibble::rownames_to_column(var = "pathway") %>%
        pivot_longer(col = -pathway, values_to = "score", names_to = "sample") %>%
        left_join(., group_infor %>% dplyr::select(sample, group))

    group_chara <- sort(unique(group_infor$group))

    if (is.null(col_used)) {
        col_name <- brewer.pal(8, "Set1")[seq_along(group_chara)]
    } else {
        col_name <- col_used[seq_along(group_chara)]
    }
    names(col_name) <- group_chara

    p <- ggplot(data_forboxplot, aes(x = pathway, y = score)) +
        geom_boxplot(aes(fill = group), notch = F, outlier.size = .5) +
        ggpubr::stat_compare_means(
            aes(pathway, score, group = group),
            label = "p.signif",
            method = ifelse(length(group_infor$group) == 2, "wilcox.test", "kruskal.test"), label.y = max(data_forboxplot$score) * 1.02
        ) +
        labs(x = "", y = "NES") +
        theme_bw(18) +
        scale_fill_manual(values = col_name) +
        theme(
            axis.text.x.bottom = element_text(vjust = 1, hjust = 1, angle = 45, colour = "black"),
            plot.margin = unit(c(1, 1, 0, 3), "cm"),
            legend.position = "top"
        )

    group_wilcoxtest_p <- map_dbl(1:nrow(ssgsea_score), function(i) {
        if (length(unique(group_infor$group)) == 2) {
            data_used <- data.frame(value = ssgsea_score[i, ] %>% as.numeric(), group = group_infor[match(colnames(ssgsea_score), group_infor$sample), "group"])
            wilcox.test_res <- wilcox.test(value ~ group, data = data_used)
            pval <- wilcox.test_res$p.value

            return(pval)
        } else {
            data_used <- data.frame(value = ssgsea_score[i, ] %>% as.numeric(), group = group_infor[match(colnames(ssgsea_score), group_infor$sample), "group"])
            kruskal.test_res <- kruskal.test(value ~ group, data = data_used)
            pval <- kruskal.test_res$p.value

            return(pval)
        }
    })

    names(group_wilcoxtest_p) <- rownames(ssgsea_score)

    p.sig_pathway <- which(group_wilcoxtest_p < 0.05) %>% names()

    library(ComplexHeatmap)

    group_chara <- sort(unique(group_infor$group))
    if (is.null(col_used)) {
        col_name <- RColorBrewer::brewer.pal(8, "Set1")[seq_along(group_chara)]
    } else {
        col_name <- col_used[seq_along(group_chara)]
    }
    names(col_name) <- group_chara

    if (is.null(cluster_name)) {
        cluster_name <- "group"
    }

    col_anno <- HeatmapAnnotation(
        df = group_infor %>% .[, "group", drop = F],
        col = list(
            group = col_name
        ),
        annotation_legend_param = list(group = list(title = cluster_name)),
        annotation_name_side = "right",
        annotation_label = cluster_name
    )

    # 显著性标注图例
    lgd <- Legend(
        labels = c("ns", "<0.05", "<0.01", "<0.001", "<0.0001"),
        title = "P.val", size = .8,
        type = "points",
        background = "white",
        pch = c(" ", "*", "**", "***", "****"),
        # border = 'black'
    )

    ssgsea_score_scaled <- t(apply(ssgsea_score, 1, function(x) scale(x, center = T, scale = T)))
    colnames(ssgsea_score_scaled) <- colnames(ssgsea_score)
    all(colnames(ssgsea_score_scaled) == group_infor$sample)

    # col_zscore <- circlize::colorRamp2(c(-1, 0, 1), c("navy", "white", "firebrick3")) #c("#0a5aa5", "white", "firebrick3")
    library(scales)
    col_zscore <- circlize::colorRamp2(c(-1, 0, 1), c("#0a5aa5", "white", "firebrick3")) # c("#0a5aa5", "white", "firebrick3")

    if (!display_all) {
        ssgsea_score_scaled <- ssgsea_score_scaled[p.sig_pathway, ]
        row_anno <- rowAnnotation(
            P.val = anno_text(
                case_when(
                    between(group_wilcoxtest_p[p.sig_pathway], 0.01, 0.05) ~ "*",
                    between(group_wilcoxtest_p[p.sig_pathway], 0.001, 0.01) ~ "**",
                    between(group_wilcoxtest_p[p.sig_pathway], 0.0001, 0.001) ~ "***",
                    group_wilcoxtest_p[p.sig_pathway] < 0.0001 ~ "****",
                    group_wilcoxtest_p[p.sig_pathway] > 0.05 ~ " "
                ),
                gp = gpar(fontsize = 10),
                location = 1,
                just = "right"
            )
        )
    } else {
        ssgsea_score_scaled <- ssgsea_score_scaled
        row_anno <- rowAnnotation(
            P.val = anno_text(
                case_when(
                    between(group_wilcoxtest_p, 0.01, 0.05) ~ "*",
                    between(group_wilcoxtest_p, 0.001, 0.01) ~ "**",
                    between(group_wilcoxtest_p, 0.0001, 0.001) ~ "***",
                    group_wilcoxtest_p < 0.0001 ~ "****",
                    group_wilcoxtest_p > 0.05 ~ " "
                ),
                gp = gpar(fontsize = 10),
                location = 1,
                just = "right"
            )
        )
    }

    cluster_res <- Heatmap(
        matrix = ssgsea_score_scaled,
        name = "NES",
        col = col_zscore,
        show_column_dend = F,
        show_column_names = F,
        show_row_names = T,
        cluster_rows = T,
        show_row_dend = F,
        # rect_gp = gpar(col = "white", lwd = 1),
        row_names_gp = gpar(fontsize = 8),
        column_title = " ",
        top_annotation = col_anno,
        row_names_max_width = max_text_width(
            rownames(ssgsea_score_scaled),
            gp = gpar(fontsize = 9.5)
        ),
        column_split = group_infor %>% .[, "group", drop = F],
        border = F,
        left_annotation = row_anno,
        row_names_side = "right"
    )

    gg.cluster_res <- cluster_res %>%
        draw(., merge_legend = T, annotation_legend_list = lgd) %>%
        grid.grabExpr() %>%
        ggplotify::as.ggplot() +
        theme(
            plot.background = element_rect(fill = "white", color = "white"),
            plot.margin = unit(c(.1, .2, .1, .1), "cm")
        )

    if (saveplot) {
        dir_now <- sprintf("%s/output/pathway_enrichment_score/", output_dir)
        if (!dir.exists(dir_now)) {
            dir.create(dir_now, recursive = T)
        } # 创建结果存储目录

        # 输出中间结果文件-----
        data.frame(
            pathway = names(group_wilcoxtest_p),
            p.val = group_wilcoxtest_p
        ) %>% write_csv(., file = sprintf("%stable_%s_score_p.csv", dir_now, var_name))

        write_csv(ssgsea_score %>% rownames_to_column(var = "pathway"), file = sprintf("%stable_%s_ssgsea_score_res.csv", dir_now, var_name))
        # 输出中间结果文件-----

        ggsave2(sprintf("%sfigure_%s_enrichment_score_boxplot.pdf", dir_now, var_name), plot = p, height = h, width = w, dpi = 300, limitsize = FALSE)
        # ggsave2(sprintf("%sfigure_%s_enrichment_score_boxplot.tiff", dir_now,var_name), plot = p, height = 9, width = nrow(ssgsea_score)*.5, dpi = 300,limitsize = FALSE)
        # ggsave2(sprintf("%sfigure_%s_enrichment_score_boxplot_dpi72.tiff", dir_now,var_name), plot = p, height = 9, width = nrow(ssgsea_score)*.5, dpi = 72,limitsize = FALSE)

        if (!display_all) {
            ggsave2(sprintf("%sfigure_%s_enrichment_score.pdf", dir_now, var_name),
                plot = gg.cluster_res,
                height = h, width = w
            )

            ggsave2(sprintf("%sfigure_%s_enrichment_score.png", dir_now, var_name),
                plot = gg.cluster_res,
                height = h, width = w, dpi = 300
            )
        } else {
            ggsave2(sprintf("%sfigure_%s_enrichment_score.pdf", dir_now, var_name),
                plot = gg.cluster_res,
                height = h, width = w
            )

            ggsave2(sprintf("%sfigure_%s_enrichment_score.png", dir_now, var_name),
                plot = gg.cluster_res,
                height = h, width = w, dpi = 300
            )
        }
    }

    res <- list(
        "ssgsea_score" = ssgsea_score,
        "wilcox_p" = group_wilcoxtest_p,
        "boxplot_p" = p,
        "heatmap_p" = gg.cluster_res
    )

    return(res)
}


# 相关性分析
# hallmark_score <- data.table::fread('/Pub/Users/wangyk/project/Poroject/GDT211207002_OV_Necroptosis/out/8.模型分组相关分析/output/pathway_enrichment_score/table_model_hallmark_ssgsea_score_res.csv')
# hallmark_score <- hallmark_score %>% as.data.frame() %>% column_to_rownames('pathway') %>% t() %>% as.data.frame()
# hallmark_score <- hallmark_score %>%
#     rownames_to_column("sample") %>%
#     inner_join(tcga_df %>% select(sample, riskscore) %>%
#         rename(Score = riskscore)) %>%
#     column_to_rownames("sample")

# colnames(hallmark_score) <- str_remove_all(colnames(hallmark_score),'HALLMARK_') %>% str_replace_all('_',' ')

# cor_res <- psych::corr.test(hallmark_score,hallmark_score,method = 'spearman',adjust = 'BH')

# pdf("/Pub/Users/wangyk/project/Poroject/GDT211207002_OV_Necroptosis/out/8.模型分组相关分析/Fig cor.pdf", width = 7, height = 7)
# p <- corrplot::corrplot(cor_res$r,
#     method = "circle",
#     type = "upper",
#     order = "original",
#     # addCoef.col = "#1C1C1C",
#     tl.col = "black",
#     tl.cex = 0.5,
#     col = colorRampPalette(c("#297eb6", "#FFFFFF","#F36C43"))(10) # ,
#     # cl.lim = c(0,1)
# )
# dev.off()
